#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/GraphTraits.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/IteratedDominanceFrontier.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Use.h"
#include "llvm/IR/Value.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/ProfileData/InstrProf.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"
#include <cassert>
#include <cstddef>
#include <utility>
#include <vector>
#include <set>
#include <queue>
#include <unordered_map>

//#define my_DEBUG_DCE 0

using namespace llvm;

namespace llvm
{
    FunctionPass *createMYDCEPass();
    void initializeMYDCEPassPass(PassRegistry &);
} // namespace llvm

namespace
{
    std::set<BasicBlock *> Dead_Block;   //存放可以到达的 Block
    std::queue<BasicBlock *> Queue;       //存放 BFS 中间结果
    std::unordered_map<BasicBlock *, ColorType> Visited;

    struct MYDCEPass : public FunctionPass
    {
        static char ID;
        MYDCEPass() : FunctionPass(ID)
        {
            initializeMYDCEPassPass(*PassRegistry::getPassRegistry());
        }

        void BFS(Function &F)   //搜索从 F 入口可达的 BasicBlock
        {
            Dead_Block.clear();
            Visited.clear();
            while (!Queue.empty())      // clear Queue
            {
                Queue.pop();
            }
            for (BasicBlock &bb : F.getBasicBlockList())
            {
                Dead_Block.insert(&bb);
                Visited[&bb] = ColorType::WHITE;
            }

            BasicBlock &Entry = F.getEntryBlock();
            Queue.push(&Entry);
            while (!Queue.empty())
            {
                BasicBlock *Cur_BB = Queue.front();
                Queue.pop();
                for (auto sSucc = succ_begin(Cur_BB), eSucc = succ_end(Cur_BB); sSucc != eSucc; ++sSucc)
                {
                    if (Visited[*sSucc] != ColorType::WHITE)   //找到其中的 BasicBlock
                    {
                        continue;
                    }
                    Visited[*sSucc] = ColorType::GRAY;
                    Queue.push(*sSucc);
                }
                Visited[Cur_BB] = ColorType::BLACK;
                Dead_Block.erase(Cur_BB);
            }
        }

        bool runOnFunction(Function &F) override
        {
            if (skipFunction(F))
            {
                return false;
            }

            SmallPtrSet<Instruction *, 32> Alive;
            SmallVector<Instruction *, 32> Worklist;

            BFS(F);

//#ifdef my_DEBUG_DCE
            llvm::outs() << "Dead Blocks: {\n";
            for (auto BB = Dead_Block.begin(); BB != Dead_Block.end(); ++BB)    //打印不会到达的基本块
            {
                llvm::outs() << (*BB)->getName() << "\n";
            }
            llvm::outs() << "}\n";

            llvm::outs() << "All Instructions: {\n";
            for (auto I = inst_begin(&F), e = inst_end(&F); I != e; ++I)
            {
                llvm::outs() << (*I) << "\n";
            }
            llvm::outs() << "}\n";
//#endif
            for (auto I = inst_begin(&F), e = inst_end(&F); I != e; ++I)
            {
                if (isa<DbgInfoIntrinsic>(*I) || (*I).isTerminator() || isa<LandingPadInst>(*I) || (*I).mayHaveSideEffects())
                {
#ifdef my_DEBUG_DCE
                    //llvm::outs() << "111" << "\n";
                    if ((*I).mayHaveSideEffects())
                    {
                        llvm::outs() << "May have side Effects: " << (*I) << "\n";
                    }
#endif
                    Alive.insert(&(*I));
                    Worklist.push_back(&(*I));
#ifdef my_DEBUG_DCE
                    llvm::outs() << (*I) << "\n";
#endif
                }
            }
            while (!Worklist.empty())   //传播生存性
            {
                Instruction *Curr = Worklist.pop_back_val();
#ifdef my_DEBUG_DCE
                //llvm::outs() << "222" << "\n";
                llvm::outs() << "Instruction: " << *Curr;
                llvm::outs() << " 's User contains: {" << "\n";
#endif
                for (Use &OI : Curr->operands())
                {
                    if (Instruction *Inst = dyn_cast<Instruction>(OI))
                    {
#ifdef my_DEBUG_DCE
                        llvm::outs() << *Inst << "\n";
#endif
                        if (Alive.insert(Inst).second)
                        {
                            Worklist.push_back(Inst);
                        }
                    }
                }
#ifdef my_DEBUG_DCE
                llvm::outs() << "}\n";
#endif
            }
#ifdef my_DEBUG_DCE
            llvm::outs() << "Instructions in Alive: {\n";
            for (auto I : Alive)
            {
                llvm::outs() << *I << "\n";
            }
            llvm::outs() << "}\n";
#endif
            llvm::outs() << "(MYDCEPass) The Eliminated Instructions: {" << "\n";
            for (auto I = inst_begin(&F), e = inst_end(&F); I != e; ++I)
            {
#ifdef my_DEBUG_DCE
                //llvm::outs() << "333" << "\n";
#endif
                if ((Dead_Block.find((*I).getParent()) != Dead_Block.end()) || !Alive.count(&(*I)))
                {
                    Worklist.push_back(&(*I));
                    llvm::outs() << (*I) << "\n";
                    (*I).dropAllReferences();
                }
            }
            llvm::outs() << "}\n";

            for (Instruction *&I : Worklist)
            {
#ifdef my_DEBUG_DCE
                //llvm::outs() << "444" << "\n";
#endif
                I->eraseFromParent();
            }
//#ifdef my_DEBUG_DCE
            llvm::outs() << "After DCE, All Instructions: {\n";
            for (auto I = inst_begin(&F), e = inst_end(&F); I != e; ++I)
            {
                llvm::outs() << (*I) << "\n";
            }
            llvm::outs() << "}\n";
//#endif
            return !Worklist.empty();
        }

        void getAnalysisUsage(AnalysisUsage &AU) const override
        {
            AU.setPreservesCFG();
        }
    };
}

char MYDCEPass::ID = 0;
INITIALIZE_PASS(MYDCEPass, "MYDCEPass", "Dead Code Eliminate", false, false)

FunctionPass * llvm::createMYDCEPass()
{
    return new MYDCEPass();
}